# Testing 'get_' methods.
# 23-10-30

import os
import sys
import jax
import jax.nn as nn
import haiku as hk
import numpy as np
from typing import Optional, Union, List

sys.path.append(os.path.dirname(sys.path[0]))

from cybertron.utils.filter import get_filter
from cybertron.utils.activation import get_activation
from cybertron.utils.cutoff import get_cutoff
from cybertron.utils.rbf import get_rbf

EPSILON = 1e-3
rng = jax.random.PRNGKey(42)

print("=================Test get_filter()====================")
def get_filter_test(x):
    # my_filter = get_filter(cls_name="dense", dim_in=4, dim_out=5, activation=nn.silu, n_hidden=2, name="my_dense_filter")
    my_filter = get_filter(cls_name="residual", dim_in=4, dim_out=5, activation=nn.silu, n_hidden=1, name="my_res_filter")

    return my_filter(x) # type: ignore

np.random.seed(42)
test_x = np.random.uniform(size=(3, 4))
filter_func = hk.transform(get_filter_test, apply_rng=True)
filter_params = filter_func.init(rng, test_x)
filter_out = jax.jit(filter_func.apply)(filter_params, rng, test_x)

print("filter_out: ", filter_out)

print("=================Test get_activation()====================")
my_act_fn = get_activation(name="silu")
in_act = jax.numpy.array([-1, 1, 2, 3, 4, 5, 1, 4], dtype=jax.numpy.float32)
print(my_act_fn(in_act))

print("=================Test get_cutoff()====================")
test_input = np.random.uniform(low=0.2, high=1.2, size=(2, 3, 3))
test_mask = test_input < 0.8

def cos_fn(x, y, training=False):
    cos = get_cutoff(cls_name="cosine", cutoff=0.8)(x, y)
    return cos

cos_fn = jax.vmap(cos_fn, in_axes=(0, 0, ))
ts_cos_fn = hk.transform(cos_fn, apply_rng=True)
params = ts_cos_fn.init(rng, test_input, test_mask)
cos_out, cos_mask_out = jax.jit(ts_cos_fn.apply)(params, rng, test_input, test_mask)
print(f"cos_out: shape of {cos_out.shape}")

print("=================Test get_rbf()====================")
# create distance mat
dist = np.random.uniform(low=0.1, high=1.2, size=(2, 3, 3))

def gs(x, training=False):
    gs_basis = get_rbf(cls_name='gaussian')(x) # type: ignore
    return gs_basis

gs = jax.vmap(gs, in_axes=(0,))
gs_fn = hk.transform(gs, apply_rng=True)
params = gs_fn.init(rng, dist)
gs_basis = jax.jit(gs_fn.apply)(params, rng, dist)
print("gs_basis: ", gs_basis.shape)
